import os
import numpy as np
import pandas as pd 


# 读取表格
dfs = []
input_dir = "./13.tree_annotation_to_table"
for index, file in enumerate(os.listdir(input_dir)):
    print(index, "\t", file)
    file_path = os.path.join(input_dir, file)
    df = pd.read_excel(file_path, keep_default_na=False)
    dfs.append(df)
df = pd.concat(dfs)

# 替换空字符串为 NaN
df['gene_local_density'] = df['gene_local_density'].replace('', np.nan)
df['gene_local_density'] = df['gene_local_density'].fillna(method='ffill')
title = df.columns.tolist()

# 读取物种相关名字或id表 
species_info_df = pd.read_csv("./cell_species.csv", header=None)
species_dict = {}
for row in species_info_df.values.tolist():
    species_dict[row[-1]] = row[0:2]

# -------------------生成 gene_info 信息表--------------------------------------
gene_info_data = []
for index, row in df.iterrows():
    try:
        genome = row["species"]
        species_id = species_dict[genome][0]
        chromosome = row["chromosome/scaffold/contif"]
        cluster_name = row["cluster_name"]
        transcript_id = row["transcript_id"]
        gene_id = row["gene_id"]
        start = row["start"]
        end = row["end"]
        strand = row["strand"]
        
        # 修改这里：正确处理空值
        density = [
            row["gene_local_density"] if pd.notna(row["gene_local_density"]) else 'Null',
            row["cluster_gene_density"] if pd.notna(row["cluster_gene_density"]) else 'Null'
        ]
        # print(density)
        
        pfam = row["pfam"]
        gene_description = row["gene_description"]
        cog_category = row["gog_category"]
        main_category = row["main_category"]
        summary_feature = row["summary_feature"]
        summary_feature = "||".join(eval(
                "['" + summary_feature.strip('"]').strip("'").strip("['").strip('"').replace(
                "', '", "---").replace('"', "'").replace("'", "").replace("---", "', '")
                + "']")
            )
        level_1 = row["level_1"]
        level_2 = row["level_2"]
        level_3 = row["level_3"]
        level_4 = row["level_4"]
        level_5 = row["level_5"]
        level_6 = row["level_6"]
        
        text = (
            [genome, species_id] + [chromosome, cluster_name, transcript_id] + 
            [gene_id, start, end, strand] + density +  
            [pfam, gene_description, cog_category, main_category, summary_feature] + 
            [level_1, level_2, level_3, level_4, level_5, level_6]
        )
        gene_info_data.append([x if pd.notna(x) else None for x in text])
    except Exception as e:
        print(e)

title = [
    "genome", "genome_id", "chromosome", "cluster_name", "transcript_id", 
    "gene_id", "start", "end", "strand", "local_gene_density", "cluster_gene_density", 
    "pfam", "gene_description", "cog_category", "main_category", "summary_feature",
    "level_1", "level_2", "level_3", "level_4", "level_5", "level_6"
]

# 输出 CSV
output_dir = "17.gene_info"
os.makedirs(output_dir, exist_ok=True)
df_out_path = os.path.join(output_dir, "gene_info.csv")
df = pd.DataFrame(gene_info_data, columns=title)
df.to_csv(df_out_path, index=False)

# 生成 SQL
def format_sql_value(x):
    if x is None:
        return "NULL"
    elif isinstance(x, (int, float)):
        return str(x)
    else:
        return f"'{str(x).replace("'", "''")}'"

sql_out_path = os.path.join(output_dir, "gene_info.sql")
sql = (
    "INSERT INTO gene_info " + str(tuple(title)).replace("'", "`") + " VALUES " + 
    ",\n".join([
        "(" + ", ".join([format_sql_value(x) for x in row]) + ")" 
        for row in gene_info_data
    ]) + ";"
)

with open(sql_out_path, "w", encoding="utf-8") as f:
    f.write(sql)
